import os.path
import time
from collections import defaultdict

import cv2
import numpy as np
import torch
from PySide6.QtCore import QThread, Signal, QMutex, QWaitCondition, QMutexLocker
from pathlib import Path

from ultralytics.cfg import get_save_dir
from ultralytics.data.utils import IMG_FORMATS, VID_FORMATS
from ultralytics.models.yolo.detect import DetectionPredictor
from ultralytics.utils import LOGGER, callbacks, MACOS, WINDOWS
from ultralytics.utils.checks import check_file, check_imshow
from ultralytics.utils.files import increment_path
from ultralytics.utils import ops
from utils import DEFAULT_CFG, get_cfg
from utils.cv_chinese import cv2AddChineseText

class PoseThread(QThread, DetectionPredictor):
    # 输入 输出 消息
    send_input = Signal(np.ndarray)
    send_output = Signal(np.ndarray)
    send_msg = Signal(str)
    # 状态栏显示数据 进度条数据
    send_fps = Signal(str)  # fps
    # send_labels = Signal(dict)  # Detected target results (number of each category)
    send_progress = Signal(int)  # Completeness
    send_class_num = Signal(int)  # Number of categories detected
    send_target_num = Signal(int)  # Targets detected

    def __init__(self, cfg=DEFAULT_CFG, overrides=None):
        super(PoseThread, self).__init__()
        self.args = get_cfg(cfg, overrides)
        self.save_dir = get_save_dir(self.args)  # 保存路径
        if self.args.show:
            self.args.show = check_imshow(warn=True)
        # 界面参数设置
        self.used_model_name = None
        self.new_model_name = None  # Models that change in real time
        self.save_res = False  # 是否保存
        self.iou_thres = 0.45  # iou
        self.conf_thres = 0.25  # conf
        self.speed_thres = 10  # delay, ms
        self.labels_dict = {}  # return a dictionary of results
        self.res_status = False  # result status
        # 线程参数设置
        self.mutex = QMutex()  # 互斥锁，用于线程同步
        self.cond = QWaitCondition()  # 条件变量，用于线程同步
        self.is_continue = True  # continue/pause
        self.max_progress_value = 0  # progress bar max value
        self.percent = 0  # progress bar percent
        # YOLO参数设置
        self.model = None
        self.data = self.args.data  # data_dict
        self.imgsz = 640
        self.dataset = None
        self.done_warmup = False
        self.nosave = False
        self.vid_cap = None
        self.name = 'exp'
        self.batch = None
        self.project = 'runs/detect'
        self.exist_ok = False
        self.callbacks = defaultdict(list, callbacks.default_callbacks)  # add callbacks
        callbacks.add_integration_callbacks(self)

        self.source = None
        self.is_file = False
        self.is_url = False
        # 网络视频流加载
        self.webcam = False
        self.screenshot = False
        self.is_folder = False
        self.quit_flag = False
        self.count = 0
        self.init_source()

        # 判断输入源是否是文件夹、文件、网络视频流

    def init_source(self):
        source = str(self.source)
        self.is_file = Path(source).suffix[1:] in IMG_FORMATS.union(VID_FORMATS)
        self.is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
        # 网络视频流加载
        self.webcam = source.isnumeric() or source.endswith('.txt') or (self.is_url and not self.is_file)
        self.screenshot = source.lower().startswith('screen')
        # 判断输入源是否是文件夹，如果是列表，则是文件夹
        self.is_folder = isinstance(self.source, list)
        if self.is_url and self.is_file:
            self.source = check_file(self.source)  # download

        # 设置当前线程为暂停状态

    def pause_thread(self):
        with QMutexLocker(self.mutex):
            self.is_continue = False

        # # 设置当前线程为继续状态

    def resume_thread(self):
        with QMutexLocker(self.mutex):
            self.is_continue = True
            self.cond.wakeOne()
        # 更新dataset

    def update_dataset(self, source):
        if self.source is not None:
            if self.is_continue:
                self.pause_thread()
            self.release_res()
            self.source = source
            self.init_source()
            self.count = 0
            self.percent = 0
            if self.is_folder:
                print("不支持文件夹输入，请输入文件")
            else:
                self.setup_source(self.source)
                self.data = iter(self.dataset)

    def run(self):
        if not self.model:
            self.send_msg.emit("Loading model: {}".format(os.path.basename(self.new_model_name)))
            self.setup_model(self.new_model_name)
            self.used_model_name = self.new_model_name
        if self.save_res:
            self.save_dir = increment_path(Path(self.project) / self.name, exist_ok=self.exist_ok)  # increment run
            (self.save_dir / 'labels').mkdir(parents=True, exist_ok=True)  # make dir
        if self.is_folder:
            for source in self.source:
                self.setup_source(source)
                self.detect()
        else:
            self.setup_source(self.source)
            self.detect()
    def detect(self, ):
        # warmup model
        if not self.done_warmup:
            self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else self.dataset.bs, 3, *self.imgsz))
            self.done_warmup = True
        self.seen, self.windows, dt, self.batch = 0, [], (ops.Profile(), ops.Profile(), ops.Profile()), None
        self.run_callbacks("on_predict_start")
        datasets = iter(self.dataset)
        count = 0
        start_time = time.time()  # used to calculate the frame rate
        while True:
            with QMutexLocker(self.mutex):
                if not self.is_continue:
                    self.cond.wait(self.mutex)
                if self.percent >= 100:
                    self.percent = 0
                if self.quit_flag:
                    return
                if self.is_continue:
                    if self.is_file:
                        self.send_msg.emit("Detecting File: {}".format(os.path.basename(self.source)))
                    elif self.webcam and not self.is_url:
                        self.send_msg.emit("Detecting Webcam: Camera_{}".format(self.source))
                    elif self.is_folder:
                        self.send_msg.emit("Detecting Folder: {}".format(os.path.dirname(self.source[0])))
                    elif self.is_url:
                        self.send_msg.emit("Detecting URL: {}".format(self.source))
                    else:
                        self.send_msg.emit("Detecting: {}".format(self.source))
                    try:
                        self.batch = next(self.data)
                    except StopIteration:
                        print("no more images")
                    else:
                        paths, im0s, s = self.batch
                        self.vid_cap = self.dataset.cap if self.dataset.mode == "video" else None
                        # 原始图片送入 input框
                        self.send_input.emit(im0s if isinstance(im0s, np.ndarray) else im0s[0])
                        n = len(im0s)
                        self.count += n
                        # 处理processBar
                        if self.vid_cap:
                            if self.vid_cap.get(cv2.CAP_PROP_FRAME_COUNT) > 0:
                                self.percent = int(
                                    self.count / self.vid_cap.get(cv2.CAP_PROP_FRAME_COUNT) * self.max_progress_value)
                                self.send_progress.emit(self.percent)
                            else:
                                self.percent = 100
                                self.send_progress.emit(self.percent)
                        else:
                            self.percent = self.max_progress_value
                        if self.count % 5 == 0 and self.count >= 5:  # Calculate the frame rate every 5 frames
                            self.send_fps.emit(str(int(5 / (time.time() - start_time))))
                            start_time = time.time()

                        # Preprocess
                        with dt[0]:
                            im = self.preprocess(im0s)
                        # Inference
                        with dt[1]:
                            preds = self.inference(im)
                        # Postprocess
                        with dt[2]:
                            self.results = self.postprocess(preds, im, im0s)

                        for i, pred in enumerate(self.results):
                            orig_img = pred.orig_img
                            self.seen += 1
                            self.results[i].speed = {
                                "preprocess": dt[0].dt * 1e3 / n,
                                "inference": dt[1].dt * 1e3 / n,
                                "postprocess": dt[2].dt * 1e3 / n,
                            }
                            if isinstance(orig_img, torch.Tensor):
                                orig_img = (orig_img.detach().permute(1, 2, 0).contiguous() * 255).to(
                                    torch.uint8).cpu().numpy()
                            if pred.boxes is not None:
                                self.deal_pred_boxes(i, pred, orig_img)

                            if self.percent == self.max_progress_value and not self.webcam:
                                self.percent = 0
                                self.send_progress.emit(0)
                                self.send_msg.emit('Finish Detection')
                                self.res_status = True
                                if self.vid_cap is not None:
                                    self.vid_cap.release()
                                print(self.vid_writer)
                                if len(self.vid_writer) > 0:
                                    for key, value in self.vid_writer.items():
                                        if isinstance(value, cv2.VideoWriter):
                                            value.release()  # release final video writer
                            # 处理结果
    def save_preds(self, vid_cap, idx, save_path):
        """Save video predictions as mp4 at specified path."""
        im0 = self.plotted_img
        suffix, fourcc = (".mp4", "avc1") if MACOS else (".avi", "WMV2") if WINDOWS else (".avi", "MJPG")
        # Save imgs
        if self.dataset.mode == "image":
            cv2.imwrite(save_path, im0)
            return save_path

        else:  # 'video' or 'stream'
            if self.vid_path[idx] != save_path:  # new video
                self.vid_path[idx] = save_path
                if isinstance(self.vid_writer[idx], cv2.VideoWriter):
                    self.vid_writer[idx].release()  # release previous video writer
                if vid_cap:  # video
                    fps = int(vid_cap.get(cv2.CAP_PROP_FPS))  # integer required, floats produce error in MP4 codec
                    w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                    h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                else:  # stream
                    fps, w, h = 30, im0.shape[1], im0.shape[0]
                # suffix, fourcc = (".mp4", "avc1") if MACOS else (".avi", "WMV2") if WINDOWS else (".avi", "MJPG")
                self.vid_writer[idx] = cv2.VideoWriter(
                    str(Path(save_path).with_suffix(suffix)), cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)
                )
            # Write video
            self.vid_writer[idx].write(im0)
            return str(Path(save_path).with_suffix(suffix))

        # 释放资源
    def deal_pred_boxes(self, index, pred, orig_img):
        target_nums = 0
        self.labels_dict = {}

        for d in reversed(pred.boxes):
            c, conf, id = int(d.cls), float(d.conf), None if d.id is None else int(d.id.item())
            name = ("" if id is None else f"id:{id} ") + pred.names[c]
            box = d.xyxy.squeeze()

            if name in self.labels_dict.keys():
                self.labels_dict[name] += 1
            else:
                self.labels_dict[name] = 1
            target_nums += 1

            cv2.rectangle(orig_img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)


            orig_img = cv2AddChineseText(orig_img, f"{name}({conf:.2f})", (int(box[0]), int(box[1])- 20),
                                         (0, 255, 0), 15)

        self.plotted_img = orig_img
        self.send_output.emit(self.plotted_img)  # after detection
        self.send_class_num.emit(len(self.labels_dict))
        self.send_target_num.emit(target_nums)
        if self.save_res:
            self.save_preds(self.vid_cap, index, self.save_dir)
        if self.speed_thres != 0:
            time.sleep(self.speed_thres / 1000)  # delay , ms


    def release_res(self):
        # 释放资源
        if self.dataset is not None:
            self.dataset.running = False  # stop flag for Thread
        # 判断self.dataset里面是否有threads
        if hasattr(self.dataset, 'threads'):
            for thread in self.dataset.threads:
                if thread.is_alive():
                    thread.join(timeout=5)  # Add timeout
        if hasattr(self.dataset, 'caps'):
            for cap in self.dataset.caps:  # Iterate through the stored VideoCapture objects
                try:
                    cap.release()  # release video capture
                except Exception as e:
                    LOGGER.warning(f"WARNING ⚠️ Could not release VideoCapture object: {e}")
        if len(self.vid_writer) > 0:
            for key, value in self.vid_writer.items():
                if isinstance(value, cv2.VideoWriter):
                    value.release()  # release final video writer
                break
        self.dataset = None  # clear dataset
        self.source = None  # input source

